#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
换手率计算器
用于计算股票的换手率，并将结果写入日线数据中
"""

import os
import json
from typing import Dict, List
import argparse
import concurrent.futures
from threading import Lock


class TurnoverRateCalculator:
    """
    换手率计算器类
    """

    def __init__(self, daydata_dir: str, hisltg_dir: str):
        """
        初始化换手率计算器
        
        :param daydata_dir: 日线数据目录路径
        :param hisltg_dir: 流通股数据目录路径
        """
        self.daydata_dir = daydata_dir
        self.hisltg_dir = hisltg_dir
        self.hisltg_data = {}
        self.processed_count = 0
        self.lock = Lock()

    def load_hisltg_for_code(self, code: str) -> None:
        """
        加载单个股票的流通股数据
        
        :param code: 股票代码
        """
        filename = f"{code}.json"
        filepath = os.path.join(self.hisltg_dir, filename)
        if os.path.exists(filepath):
            with open(filepath, 'r', encoding='utf-8') as f:
                self.hisltg_data[code] = json.load(f)

    def get_free_float_for_date(self, code: str, date: int) -> float:
        """
        获取指定股票在指定日期的流通股数
        
        :param code: 股票代码
        :param date: 日期（时间戳）
        :return: 流通股数（万股），如果找不到则返回None
        """
        if code not in self.hisltg_data:
            return None
            
        # 查找最接近且不大于该日期的流通股数据
        free_float = None
        for item in self.hisltg_data[code]:
            if item['change_date'] <= date:
                free_float = item['free_float']
            else:
                break
                
        return free_float

    def calculate_turnover_rate(self, volume: int, free_float: float) -> float:
        """
        计算换手率
        
        换手率 = 成交量 / 流通股数
        
        :param volume: 成交量（股）
        :param free_float: 流通股数（万股）
        :return: 换手率（百分比）
        """
        if free_float is None or free_float <= 0:
            return None
            
        # 注意单位转换：
        # volume 是股，需要转成万股
        # free_float 已经是万股
        turnover_rate = (volume * 0.0001) / free_float * 100
        return round(turnover_rate, 4)

    def process_single_stock(self, filename: str, output_dir: str = None) -> bool:
        """
        处理单个股票的日线数据并计算换手率
        
        :param filename: 日线数据文件名
        :param output_dir: 输出目录路径，如果为None则覆盖原文件
        :return: 处理是否成功
        """
        try:
            if not filename.endswith('.json'):
                return False
                
            code = filename[:-5]  # 移除 .json 后缀
            
            # 加载该股票的流通股数据
            self.load_hisltg_for_code(code)
            
            # 加载日线数据
            filepath = os.path.join(self.daydata_dir, filename)
            with open(filepath, 'r', encoding='utf-8') as f:
                daydata_list = json.load(f)
            
            # 处理每条日线数据
            updated_count = 0
            for daydata in daydata_list:
                date = int(daydata['date'])
                volume = daydata['volume']
                
                # 获取该日期的流通股数
                free_float = self.get_free_float_for_date(code, date)
                
                # 计算换手率
                turnover_rate = self.calculate_turnover_rate(volume, free_float)
                
                # 添加换手率到日线数据中
                daydata['turnover_rate'] = turnover_rate
                if turnover_rate is not None:
                    updated_count += 1
            
            # 保存更新后的数据
            if output_dir:
                output_filepath = os.path.join(output_dir, filename)
            else:
                output_filepath = filepath
                
            with open(output_filepath, 'w', encoding='utf-8') as f:
                json.dump(daydata_list, f, ensure_ascii=False, indent=2)
            
            # 清理该股票的流通股数据以节省内存
            if code in self.hisltg_data:
                del self.hisltg_data[code]
            
            # 更新计数器
            with self.lock:
                self.processed_count += 1
                print(f"[{self.processed_count}] 股票 {code} 处理完成，成功计算 {updated_count}/{len(daydata_list)} 条记录的换手率")
            
            return True
        except Exception as e:
            print(f"处理文件 {filename} 时发生错误: {e}")
            return False

    def run(self, output_dir: str = None, max_workers: int = 4) -> None:
        """
        运行换手率计算主流程
        
        :param output_dir: 输出目录路径，如果为None则覆盖原文件
        :param max_workers: 最大线程数
        """
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
        
        # 获取所有日线数据文件
        daydata_files = [f for f in os.listdir(self.daydata_dir) if f.endswith('.json')]
        print(f"发现 {len(daydata_files)} 个日线数据文件待处理")
        
        # 重置计数器
        self.processed_count = 0
        
        # 使用线程池处理文件
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = []
            for filename in daydata_files:
                future = executor.submit(self.process_single_stock, filename, output_dir)
                futures.append(future)
            
            # 等待所有任务完成
            concurrent.futures.wait(futures)
            
        # 统计结果
        success_count = sum(1 for future in futures if future.result())
        print(f"处理完成，成功处理 {success_count}/{len(daydata_files)} 个文件")


def main():
    """
    主函数
    """
    parser = argparse.ArgumentParser(description='换手率计算器')
    parser.add_argument('daydata_dir', help='日线数据目录路径')
    parser.add_argument('hisltg_dir', help='流通股数据目录路径')
    parser.add_argument('-o', '--output', help='输出目录路径（可选，默认覆盖原文件）')
    parser.add_argument('-w', '--workers', type=int, default=4, help='最大线程数，默认为4')
    
    args = parser.parse_args()
    
    calculator = TurnoverRateCalculator(args.daydata_dir, args.hisltg_dir)
    calculator.run(args.output, args.workers)


if __name__ == '__main__':
    main()
    